package com.tuling.llmlangchain4j.storememery;

import com.mongodb.client.result.UpdateResult;
import com.tuling.llmlangchain4j.bean.ChatMessages;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.LinkedList;
import java.util.List;
@Component
public class MemerryStorePersistent implements ChatMemoryStore {


    @Autowired
    private MongoTemplate mongoTemplate;


    public List<ChatMessage> getMessages(Object memoryId) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);
        if(chatMessages == null){
            return new LinkedList<>();
        }
        String content = chatMessages.getContent();
        return ChatMessageDeserializer.messagesFromJson(content);
    }

    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        String json = ChatMessageSerializer.messagesToJson(messages);
        Update updateMessage = new Update().set("content",json);
        UpdateResult upsert = mongoTemplate.upsert(query, updateMessage, ChatMessages.class);
    }

    public void deleteMessages(Object memoryId) {
        Criteria criteriId = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteriId);
        mongoTemplate.remove(query, ChatMessages.class);
    }


}
